##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Yaoyao Liu
## NUS School of Computing
## Email: yaoyao.liu@nus.edu.sg
## Copyright (c) 2019
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import argparse
import tempfile
import tarfile
import os
import csv
import cv2
import shutil
from tqdm import tqdm
import glob

# argument parser
parser = argparse.ArgumentParser(description='')
parser.add_argument('--tar_path', type=str)
parser.add_argument('--target_dir', type=str)
parser.add_argument('--image_resize', type=int, default=84)

args = parser.parse_args()

titles = ["train", "val", "test"]
current_path = os.path.split(os.path.realpath(__file__))[0]
csv_paths = [os.path.join(current_path, 'csv_files', csv) for csv in ['train.csv', 'val.csv', 'test.csv']]

class_names = ['n02110341', 'n01930112', 'n04509417', 'n04067472', 'n04515003', 'n02120079', 'n03924679',
               'n02687172', 'n03075370', 'n07747607', 'n09246464', 'n02457408', 'n04418357', 'n03535780',
               'n04435653', 'n03207743', 'n04251144', 'n03062245', 'n02174001', 'n07613480', 'n03998194',
               'n02074367', 'n04146614', 'n04243546', 'n03854065', 'n03838899', 'n02871525', 'n03544143',
               'n02108089', 'n13133613', 'n03676483', 'n03337140', 'n03272010', 'n01770081', 'n09256479',
               'n02091244', 'n02116738', 'n04275548', 'n03773504', 'n02606052', 'n03146219', 'n04149813',
               'n07697537', 'n02823428', 'n02089867', 'n03017168', 'n01704323', 'n01532829', 'n03047690',
               'n03775546', 'n01843383', 'n02971356', 'n13054560', 'n02108551', 'n02101006', 'n03417042',
               'n04612504', 'n01558993', 'n04522168', 'n02795169', 'n06794110', 'n01855672', 'n04258138',
               'n02110063', 'n07584110', 'n02091831', 'n03584254', 'n03888605', 'n02113712', 'n03980874',
               'n02219486', 'n02138441', 'n02165456', 'n02108915', 'n03770439', 'n01981276', 'n03220513',
               'n02099601', 'n02747177', 'n01749939', 'n03476684', 'n02105505', 'n02950826', 'n04389033',
               'n03347037', 'n02966193', 'n03127925', 'n03400231', 'n04296562', 'n03527444', 'n04443257',
               'n02443484', 'n02114548', 'n04604644', 'n01910747', 'n04596742', 'n02111277', 'n03908618',
               'n02129165', 'n02981792']


class MiniImageNetGenerator(object):
    def __init__(self, input_args):
        self.image_resize = input_args.image_resize
        self.imagenet_dir = tempfile.mkdtemp()
        self.target_dir = os.path.expanduser(input_args.target_dir)

        if input_args.tar_path is None:
            raise ValueError("tar_path cannot be none, path of ILSVRC2012_img_train.tar must be given")

        print('')
        print('Extracting images from ILSVRC2012_img_train.tar !!! \n')
        tarfile.open(input_args.tar_path).extractall(self.imagenet_dir)
        print('ILSVRC2012_img_train.tar has been already extracted !!! \n\n')

        print("Extracting images from sub tar files ......")
        process_bar = tqdm(class_names)
        for class_name in process_bar:
            process_bar.set_description("Extracting %s" % (class_name + ".tar"))

            sub_tar_path = os.path.join(self.imagenet_dir, class_name + '.tar')
            if not os.path.exists(sub_tar_path):
                raise ValueError("{} cannot be found in {}".format(class_name + ".tar", self.imagenet_dir))

            sub_class_dir = os.path.join(self.imagenet_dir, class_name)
            if not os.path.exists(sub_class_dir):
                os.makedirs(sub_class_dir)
            else:
                print("Warning: {} has already existed !!!".format(sub_class_dir))
                shutil.rmtree(sub_class_dir)
                print("Warning: {} has already removed !!!".format(sub_class_dir))

            # print('extracting images from {} !!!'.format(sub_tar_path))
            tarfile.open(sub_tar_path).extractall(sub_class_dir)
            # print('{} has been already extracted !!!'.format(sub_tar_path))

    def process_original_files(self):
        if not os.path.exists(self.target_dir):
            os.makedirs(self.target_dir)

        processed_img_dir = os.path.join(self.target_dir, 'processed_images')
        if os.path.exists(processed_img_dir):
            print("Warning: {} has already existed !!!".format(processed_img_dir))
            shutil.rmtree(processed_img_dir)
            print("Warning: {} has already removed !!!".format(processed_img_dir))
        os.makedirs(processed_img_dir)

        processed_split_dirs = []
        for title in titles:  # titles = ["train", "val", "test"]
            target_split_dir = os.path.join(processed_img_dir, title)
            processed_split_dirs.append(target_split_dir)
            os.makedirs(target_split_dir)

        for csv_path in csv_paths:
            if not os.path.exists(csv_path):
                raise ValueError("Error: {} cannot be found !!!".format(csv_path))

        for target_split_dir, csv_path in zip(processed_split_dirs, csv_paths):
            with open(csv_path) as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                next(csv_reader, None)
                images_csv = {}
                print('\n\n')
                print('Reading Images according with {} ...... '.format(csv_path))

                for row in csv_reader:
                    img_name = row[0]
                    img_Ids = int(img_name[img_name.index('.') - 4: img_name.index('.')])
                    if row[1] in images_csv.keys():
                        images_csv[row[1]].append((img_Ids - 1, img_name))
                    else:
                        images_csv[row[1]] = [(img_Ids - 1, img_name)]

                print('Writing photos to {} ...... '.format(target_split_dir))

                process_bar = tqdm(images_csv.keys())
                for class_name in process_bar:
                    process_bar.set_description("Processing %s" % class_name)

                    source_dir = os.path.join(self.imagenet_dir, class_name)
                    target_dir = os.path.join(target_split_dir, class_name)
                    if os.path.exists(target_dir):
                        print("Warning: {} has already existed !!!".format(target_dir))
                        shutil.rmtree(target_dir)
                        print("Warning: {} has already removed !!!".format(target_dir))
                    os.makedirs(target_dir)

                    source_files = []
                    for file in glob.glob(os.path.join(source_dir, "*" + class_name + "*")):
                        source_files.append(file)
                    source_files_Ids = [int(img_path[img_path.rfind('_') + 1:img_path.rfind('.')]) for img_path in
                                        source_files]
                    source_files_index_sorted = sorted(range(len(source_files)), key=source_files_Ids.__getitem__)

                    for img_Ids, img_name in images_csv[class_name]:
                        source_file = source_files[source_files_index_sorted[img_Ids]]
                        target_file = os.path.join(target_dir, img_name)

                        if not os.path.exists(source_file):
                            raise ValueError("Error: {} cannot be found !!!".format(source_file))

                        if self.image_resize == 0:
                            shutil.copyfile(source_file, target_file)
                        else:
                            im = cv2.imread(source_file)
                            im_resized = cv2.resize(im, (self.image_resize, self.image_resize),
                                                    interpolation=cv2.INTER_AREA)
                            cv2.imwrite(target_file, im_resized)

    def clean(self):
        shutil.rmtree(self.imagenet_dir)
        print('')


if __name__ == "__main__":
    mini_imagenet_dataset = MiniImageNetGenerator(args)

    mini_imagenet_dataset.process_original_files()

    mini_imagenet_dataset.clean()
